# -*- coding: utf-8 -*-
# @Time    : 2020/6/20 下午1:27
# @Author  : caotian
# @FileName: lenettrain2.py
# @Software: PyCharm
import numpy as np
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
from PIL import Image
import os
import sys
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import lenetmodel as lm

def train(model):
    print("start training.....")
    model.train()
    epochnum=5
    opt=fluid.optimizer.Momentum(learning_rate=0.001,momentum=0.9,parameter_list=model.parameters())
    train_loader=paddle.batch(paddle.dataset.mnist.train(),batch_size=10)
    valid_loader=paddle.batch(paddle.dataset.mnist.test(),batch_size=10)
    for epoch in range(epochnum):
        for batchid,data in enumerate(train_loader()):
            xdata=np.array([item[0] for item in data],dtype='float32').reshape(-1,1,28,28)
            ydata=np.array([item[1] for item in data],dtype='int64').reshape(-1,1)
            img=fluid.dygraph.to_variable(xdata)
            label=fluid.dygraph.to_variable(ydata)
            logits=model(img)
            loss=fluid.layers.softmax_with_cross_entropy(logits,label)
            avg_loss=fluid.layers.mean(loss)
            if batchid % 1000 == 0:
                print("epoch: {}, batch_id: {}, loss is: {}".format(epoch, batchid, avg_loss.numpy()))
            avg_loss.backward()
            opt.minimize(avg_loss)
            model.clear_gradients()
        model.eval()
        accuracies=[]
        losses=[]
        for batch_id,data in enumerate(valid_loader()):
            xdata=np.array([item[0] for item in data],dtype='float32').reshape(-1,1,28,28)
            ydata=np.array([item[1] for item in data],dtype='int64').reshape(-1,1)
            img=fluid.dygraph.to_variable(xdata)
            label=fluid.dygraph.to_variable(ydata)
            logits=model(img)
            pred=fluid.layers.softmax(logits)
            loss=fluid.layers.softmax_with_cross_entropy(logits,label)
            acc=fluid.layers.accuracy(pred,label)
            accuracies.append(acc.numpy())
            losses.append(loss.numpy())
        print("[validation accuracy/loss:{}/{}".format(np.mean(accuracies),np.mean(losses)))
        model.train()
    fluid.save_dygraph(model.state_dict(),'lenetmnist-model')

if __name__ == '__main__':
    with fluid.dygraph.guard():
        model=lm.LeNet(num_classes=10)
        train(model)



